import tiktoken
import torch  # PyTouch 深度学习框架
from torch.utils.data import Dataset, DataLoader  # PyTorch 数据加载工具


class GPTDataSetV1(Dataset):
    """自定义 PyTorch Dataset，用于处理 GPT 风格的训练数据"""

    def __init__(self, txt, tokenizer, max_length, stride):
        """
        初始化数据集：
        - txt: 输入的原始文本
        - tokenizer: 分词器（这里使用 tiktoken）
        - max_length: 每个输入序列的最大长度
        - stride: 滑动窗口的步长（控制数据重叠程度）
        """
        self.input_ids = []  # 存储输入序列（token IDs）
        self.target_ids = []  # 存储目标序列（比输入序列右移 1 个 token）

        # 1. 对整个文本进行分词，得到 token IDs（例如 [123, 456, 789, ...]）
        token_ids = tokenizer.encode(txt)
        # 2. 使用滑动窗口生成训练样本（input 和 target）
        for i in range(0, len(token_ids) - max_length, stride):
            # 当前窗口的输入（从 i 到 i+max_length）
            input_chunk = token_ids[i: i + max_length]
            # 当前窗口的目标（从 i+1 到 i+max_length+1，即输入右移 1 位）
            target_chunk = token_ids[i + 1: i + max_length + 1]
            # 将输入和目标转换为 PyTorch Tensor 并存储
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        """根据索引返回单个样本（input 和 target）"""
        return self.input_ids[idx], self.target_ids[idx]


def create_dataloader_v1(txt, batch_size=4, max_length=256,
                         stride=128, shuffle=True, drop_last=True,
                         num_workers=0):
    """
    创建 PyTorch DataLoader，用于批量加载数据：
    - txt: 输入的原始文本
    - batch_size: 每批次的样本数
    - max_length: 输入序列的最大长度
    - stride: 滑动窗口的步长
    - shuffle: 是否打乱数据顺序
    - drop_last: 是否丢弃最后一个不完整的批次
    - num_workers: 数据加载的线程数（0 表示主线程加载）
    """
    # 1. 初始化分词器（使用 GPT-2 的分词方式）
    tokenizer = tiktoken.get_encoding('gpt2')

    # 2. 创建自定义 Dataset
    dataset = GPTDataSetV1(txt, tokenizer, max_length, stride)

    # 3. 创建 DataLoader，用于批量加载数据
    dataloader = DataLoader(
        dataset=dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=drop_last,
        num_workers=num_workers)
    return dataloader


with open('../data/the-verdict.txt', 'r', encoding='utf-8') as f:
    raw_text = f.read()

dataloader = create_dataloader_v1(raw_text, batch_size=1, max_length=4, stride=1, shuffle=False)
# 获取第一个批次的数据并打印
data_iter = iter(dataloader)
first_batch = next(data_iter)
print(first_batch)

second_batch = next(data_iter)
print(second_batch)

# 批次大于1的使用数据加载器采样
dataloader = create_dataloader_v1(raw_text, batch_size=8, max_length=4, stride=4, shuffle=False)
data_iter = iter(dataloader)
inputs, targets = next(data_iter)
print("Inputs:\n", inputs)
print("\nTargets:\n", targets)

# 值得说明的是，我们选择将步幅增加到4来充分利用数据集（不会跳过任何一个单词），同时避免不同批次之间的数据重叠，因为过多的重叠可能会增加模型过拟合的风险。
